from typing import Any, Optional, Union
import jax.numpy as jnp

Array = Any


def softmax1(x: Array, axis: Optional[Union[int, tuple[int, ...]]] = -1):
    # x_max = jnp.max(x, axis, keepdims=True)
    # unnormalized = jnp.exp(x - x_max)
    # result = unnormalized / (1 + jnp.sum(unnormalized, axis, keepdims=True))
    unnormalized = jnp.exp(x)
    result = unnormalized / (1 + jnp.sum(unnormalized, axis, keepdims=True))
    return result

def _softmax1(x: Array, axis: Optional[Union[int, tuple[int, ...]]] = -1):
    x_max = jnp.max(x, axis, keepdims=True)
    unnormalized = jnp.exp(x - x_max)
    result = unnormalized / (1 + jnp.sum(unnormalized, axis, keepdims=True))
    return result
